{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "25ed24e6",
   "metadata": {},
   "source": [
    "## Bert Layer Demonstration\n",
    "The purpose of this notebook is to demonstrate the computation of a single BERT layer with the least complexity and verify it's correctness with a reference implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e6f73243",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "06fe61e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BertEncoder(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])\n",
    "        self.gradient_checkpointing = False\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        hidden_states,\n",
    "        attention_mask=None,\n",
    "        head_mask=None,\n",
    "        encoder_hidden_states=None,\n",
    "        encoder_attention_mask=None,\n",
    "        past_key_values=None,\n",
    "        use_cache=None,\n",
    "        output_attentions=False,\n",
    "        output_hidden_states=False,\n",
    "        return_dict=True,\n",
    "    ):\n",
    "        all_hidden_states = () if output_hidden_states else None\n",
    "        all_self_attentions = () if output_attentions else None\n",
    "        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None\n",
    "\n",
    "        next_decoder_cache = () if use_cache else None\n",
    "        for i, layer_module in enumerate(self.layer):\n",
    "            if output_hidden_states:\n",
    "                all_hidden_states = all_hidden_states + (hidden_states,)\n",
    "\n",
    "            layer_head_mask = head_mask[i] if head_mask is not None else None\n",
    "            past_key_value = past_key_values[i] if past_key_values is not None else None\n",
    "\n",
    "            if self.gradient_checkpointing and self.training:\n",
    "\n",
    "                if use_cache:\n",
    "                    logger.warning(\n",
    "                        \"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\"\n",
    "                    )\n",
    "                    use_cache = False\n",
    "\n",
    "                def create_custom_forward(module):\n",
    "                    def custom_forward(*inputs):\n",
    "                        return module(*inputs, past_key_value, output_attentions)\n",
    "\n",
    "                    return custom_forward\n",
    "\n",
    "                layer_outputs = torch.utils.checkpoint.checkpoint(\n",
    "                    create_custom_forward(layer_module),\n",
    "                    hidden_states,\n",
    "                    attention_mask,\n",
    "                    layer_head_mask,\n",
    "                    encoder_hidden_states,\n",
    "                    encoder_attention_mask,\n",
    "                )\n",
    "            else:\n",
    "                layer_outputs = layer_module(\n",
    "                    hidden_states,\n",
    "                    attention_mask,\n",
    "                    layer_head_mask,\n",
    "                    encoder_hidden_states,\n",
    "                    encoder_attention_mask,\n",
    "                    past_key_value,\n",
    "                    output_attentions,\n",
    "                )\n",
    "\n",
    "            hidden_states = layer_outputs[0]\n",
    "            if use_cache:\n",
    "                next_decoder_cache += (layer_outputs[-1],)\n",
    "            if output_attentions:\n",
    "                all_self_attentions = all_self_attentions + (layer_outputs[1],)\n",
    "                if self.config.add_cross_attention:\n",
    "                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)\n",
    "\n",
    "        if output_hidden_states:\n",
    "            all_hidden_states = all_hidden_states + (hidden_states,)\n",
    "\n",
    "        if not return_dict:\n",
    "            return tuple(\n",
    "                v\n",
    "                for v in [\n",
    "                    hidden_states,\n",
    "                    next_decoder_cache,\n",
    "                    all_hidden_states,\n",
    "                    all_self_attentions,\n",
    "                    all_cross_attentions,\n",
    "                ]\n",
    "                if v is not None\n",
    "            )\n",
    "        return BaseModelOutputWithPastAndCrossAttentions(\n",
    "            last_hidden_state=hidden_states,\n",
    "            past_key_values=next_decoder_cache,\n",
    "            hidden_states=all_hidden_states,\n",
    "            attentions=all_self_attentions,\n",
    "            cross_attentions=all_cross_attentions,\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "81ef0863",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BertLayer(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.chunk_size_feed_forward = config.chunk_size_feed_forward\n",
    "        self.seq_len_dim = 1\n",
    "        self.attention = BertAttention(config)\n",
    "        self.is_decoder = config.is_decoder\n",
    "        self.add_cross_attention = config.add_cross_attention\n",
    "        if self.add_cross_attention:\n",
    "            if not self.is_decoder:\n",
    "                raise ValueError(f\"{self} should be used as a decoder model if cross attention is added\")\n",
    "            self.crossattention = BertAttention(config, position_embedding_type=\"absolute\")\n",
    "        self.intermediate = BertIntermediate(config)\n",
    "        self.output = BertOutput(config)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        hidden_states,\n",
    "        attention_mask=None,\n",
    "        head_mask=None,\n",
    "        encoder_hidden_states=None,\n",
    "        encoder_attention_mask=None,\n",
    "        past_key_value=None,\n",
    "        output_attentions=False,\n",
    "    ):\n",
    "        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2\n",
    "        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None\n",
    "        self_attention_outputs = self.attention(\n",
    "            hidden_states,\n",
    "            attention_mask,\n",
    "            head_mask,\n",
    "            output_attentions=output_attentions,\n",
    "            past_key_value=self_attn_past_key_value,\n",
    "        )\n",
    "        attention_output = self_attention_outputs[0]\n",
    "\n",
    "        # if decoder, the last output is tuple of self-attn cache\n",
    "        if self.is_decoder:\n",
    "            outputs = self_attention_outputs[1:-1]\n",
    "            present_key_value = self_attention_outputs[-1]\n",
    "        else:\n",
    "            outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights\n",
    "\n",
    "        cross_attn_present_key_value = None\n",
    "        if self.is_decoder and encoder_hidden_states is not None:\n",
    "            if not hasattr(self, \"crossattention\"):\n",
    "                raise ValueError(\n",
    "                    f\"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`\"\n",
    "                )\n",
    "\n",
    "            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple\n",
    "            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None\n",
    "            cross_attention_outputs = self.crossattention(\n",
    "                attention_output,\n",
    "                attention_mask,\n",
    "                head_mask,\n",
    "                encoder_hidden_states,\n",
    "                encoder_attention_mask,\n",
    "                cross_attn_past_key_value,\n",
    "                output_attentions,\n",
    "            )\n",
    "            attention_output = cross_attention_outputs[0]\n",
    "            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights\n",
    "\n",
    "            # add cross-attn cache to positions 3,4 of present_key_value tuple\n",
    "            cross_attn_present_key_value = cross_attention_outputs[-1]\n",
    "            present_key_value = present_key_value + cross_attn_present_key_value\n",
    "\n",
    "        layer_output = apply_chunking_to_forward(\n",
    "            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output\n",
    "        )\n",
    "        outputs = (layer_output,) + outputs\n",
    "\n",
    "        # if decoder, return the attn key/values as the last output\n",
    "        if self.is_decoder:\n",
    "            outputs = outputs + (present_key_value,)\n",
    "\n",
    "        return outputs\n",
    "\n",
    "    def feed_forward_chunk(self, attention_output):\n",
    "        intermediate_output = self.intermediate(attention_output)\n",
    "        layer_output = self.output(intermediate_output, attention_output)\n",
    "        return layer_output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dcf2957e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BertAttention(nn.Module):\n",
    "    def __init__(self, config, position_embedding_type=None):\n",
    "        super().__init__()\n",
    "        self.self = BertSelfAttention(config, position_embedding_type=position_embedding_type)\n",
    "        self.output = BertSelfOutput(config)\n",
    "        self.pruned_heads = set()\n",
    "\n",
    "    def prune_heads(self, heads):\n",
    "        if len(heads) == 0:\n",
    "            return\n",
    "        heads, index = find_pruneable_heads_and_indices(\n",
    "            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads\n",
    "        )\n",
    "\n",
    "        # Prune linear layers\n",
    "        self.self.query = prune_linear_layer(self.self.query, index)\n",
    "        self.self.key = prune_linear_layer(self.self.key, index)\n",
    "        self.self.value = prune_linear_layer(self.self.value, index)\n",
    "        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)\n",
    "\n",
    "        # Update hyper params and store pruned heads\n",
    "        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)\n",
    "        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads\n",
    "        self.pruned_heads = self.pruned_heads.union(heads)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        hidden_states,\n",
    "        attention_mask=None,\n",
    "        head_mask=None,\n",
    "        encoder_hidden_states=None,\n",
    "        encoder_attention_mask=None,\n",
    "        past_key_value=None,\n",
    "        output_attentions=False,\n",
    "    ):\n",
    "        self_outputs = self.self(\n",
    "            hidden_states,\n",
    "            attention_mask,\n",
    "            head_mask,\n",
    "            encoder_hidden_states,\n",
    "            encoder_attention_mask,\n",
    "            past_key_value,\n",
    "            output_attentions,\n",
    "        )\n",
    "        attention_output = self.output(self_outputs[0], hidden_states)\n",
    "        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them\n",
    "        return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d4325919",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BertSelfOutput(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.dense = nn.Linear(config.hidden_size, config.hidden_size)\n",
    "        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n",
    "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
    "\n",
    "    def forward(self, hidden_states, input_tensor):\n",
    "        hidden_states = self.dense(hidden_states)\n",
    "        hidden_states = self.dropout(hidden_states)\n",
    "        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n",
    "        return hidden_states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9e781aa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BertSelfAttention(nn.Module):\n",
    "    def __init__(self, config, position_embedding_type=None):\n",
    "        super().__init__()\n",
    "        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, \"embedding_size\"):\n",
    "            raise ValueError(\n",
    "                f\"The hidden size ({config.hidden_size}) is not a multiple of the number of attention \"\n",
    "                f\"heads ({config.num_attention_heads})\"\n",
    "            )\n",
    "\n",
    "        self.num_attention_heads = config.num_attention_heads\n",
    "        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)\n",
    "        self.all_head_size = self.num_attention_heads * self.attention_head_size\n",
    "\n",
    "        self.query = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "        self.key = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "        self.value = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "\n",
    "        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)\n",
    "        self.position_embedding_type = position_embedding_type or getattr(\n",
    "            config, \"position_embedding_type\", \"absolute\"\n",
    "        )\n",
    "        if self.position_embedding_type == \"relative_key\" or self.position_embedding_type == \"relative_key_query\":\n",
    "            self.max_position_embeddings = config.max_position_embeddings\n",
    "            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)\n",
    "\n",
    "        self.is_decoder = config.is_decoder\n",
    "\n",
    "    def transpose_for_scores(self, x):\n",
    "        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)\n",
    "        x = x.view(new_x_shape)\n",
    "        return x.permute(0, 2, 1, 3)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        hidden_states,\n",
    "        attention_mask=None,\n",
    "        head_mask=None,\n",
    "        encoder_hidden_states=None,\n",
    "        encoder_attention_mask=None,\n",
    "        past_key_value=None,\n",
    "        output_attentions=False,\n",
    "    ):\n",
    "        mixed_query_layer = self.query(hidden_states)\n",
    "\n",
    "        # If this is instantiated as a cross-attention module, the keys\n",
    "        # and values come from an encoder; the attention mask needs to be\n",
    "        # such that the encoder's padding tokens are not attended to.\n",
    "        is_cross_attention = encoder_hidden_states is not None\n",
    "\n",
    "        if is_cross_attention and past_key_value is not None:\n",
    "            # reuse k,v, cross_attentions\n",
    "            key_layer = past_key_value[0]\n",
    "            value_layer = past_key_value[1]\n",
    "            attention_mask = encoder_attention_mask\n",
    "        elif is_cross_attention:\n",
    "            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))\n",
    "            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))\n",
    "            attention_mask = encoder_attention_mask\n",
    "        elif past_key_value is not None:\n",
    "            key_layer = self.transpose_for_scores(self.key(hidden_states))\n",
    "            value_layer = self.transpose_for_scores(self.value(hidden_states))\n",
    "            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)\n",
    "            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)\n",
    "        else:\n",
    "            key_layer = self.transpose_for_scores(self.key(hidden_states))\n",
    "            value_layer = self.transpose_for_scores(self.value(hidden_states))\n",
    "\n",
    "        query_layer = self.transpose_for_scores(mixed_query_layer)\n",
    "\n",
    "        if self.is_decoder:\n",
    "            # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.\n",
    "            # Further calls to cross_attention layer can then reuse all cross-attention\n",
    "            # key/value_states (first \"if\" case)\n",
    "            # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of\n",
    "            # all previous decoder key/value_states. Further calls to uni-directional self-attention\n",
    "            # can concat previous decoder key/value_states to current projected key/value_states (third \"elif\" case)\n",
    "            # if encoder bi-directional self-attention `past_key_value` is always `None`\n",
    "            past_key_value = (key_layer, value_layer)\n",
    "\n",
    "        # Take the dot product between \"query\" and \"key\" to get the raw attention scores.\n",
    "        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))\n",
    "\n",
    "        if self.position_embedding_type == \"relative_key\" or self.position_embedding_type == \"relative_key_query\":\n",
    "            seq_length = hidden_states.size()[1]\n",
    "            position_ids_l = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)\n",
    "            position_ids_r = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device).view(1, -1)\n",
    "            distance = position_ids_l - position_ids_r\n",
    "            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)\n",
    "            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  # fp16 compatibility\n",
    "\n",
    "            if self.position_embedding_type == \"relative_key\":\n",
    "                relative_position_scores = torch.einsum(\"bhld,lrd->bhlr\", query_layer, positional_embedding)\n",
    "                attention_scores = attention_scores + relative_position_scores\n",
    "            elif self.position_embedding_type == \"relative_key_query\":\n",
    "                relative_position_scores_query = torch.einsum(\"bhld,lrd->bhlr\", query_layer, positional_embedding)\n",
    "                relative_position_scores_key = torch.einsum(\"bhrd,lrd->bhlr\", key_layer, positional_embedding)\n",
    "                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key\n",
    "\n",
    "        attention_scores = attention_scores / math.sqrt(self.attention_head_size)\n",
    "        if attention_mask is not None:\n",
    "            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)\n",
    "            attention_scores = attention_scores + attention_mask\n",
    "\n",
    "        # Normalize the attention scores to probabilities.\n",
    "        attention_probs = nn.functional.softmax(attention_scores, dim=-1)\n",
    "\n",
    "        # This is actually dropping out entire tokens to attend to, which might\n",
    "        # seem a bit unusual, but is taken from the original Transformer paper.\n",
    "        attention_probs = self.dropout(attention_probs)\n",
    "\n",
    "        # Mask heads if we want to\n",
    "        if head_mask is not None:\n",
    "            attention_probs = attention_probs * head_mask\n",
    "\n",
    "        context_layer = torch.matmul(attention_probs, value_layer)\n",
    "\n",
    "        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()\n",
    "        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)\n",
    "        context_layer = context_layer.view(new_context_layer_shape)\n",
    "\n",
    "        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)\n",
    "\n",
    "        if self.is_decoder:\n",
    "            outputs = outputs + (past_key_value,)\n",
    "        return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9730a0c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BertIntermediate(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)\n",
    "        if isinstance(config.hidden_act, str):\n",
    "            self.intermediate_act_fn = ACT2FN[config.hidden_act]\n",
    "        else:\n",
    "            self.intermediate_act_fn = config.hidden_act\n",
    "\n",
    "    def forward(self, hidden_states):\n",
    "        hidden_states = self.dense(hidden_states)\n",
    "        hidden_states = self.intermediate_act_fn(hidden_states)\n",
    "        return hidden_states\n",
    "\n",
    "\n",
    "class BertOutput(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)\n",
    "        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)\n",
    "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
    "\n",
    "    def forward(self, hidden_states, input_tensor):\n",
    "        hidden_states = self.dense(hidden_states)\n",
    "        hidden_states = self.dropout(hidden_states)\n",
    "        hidden_states = self.LayerNorm(hidden_states + input_tensor)\n",
    "        return hidden_states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "095d0f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4ac0d36",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88192fe6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b58ec12d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98a17d8f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1251ac30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import BertTokenizer, BertModel\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
    "model = BertModel.from_pretrained(\"bert-base-uncased\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2be8b61b",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"I would love to get a degree!\"\n",
    "encoded_input = tokenizer(text, return_tensors='pt')\n",
    "encoded_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83f1662b",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = model(**encoded_input)\n",
    "output"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "trans-fat",
   "language": "python",
   "name": "trans-fat"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
